#!/bin/bash

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 函数:打印颜色文本
print_color() {
    local color=$1
    local text=$2
    echo -e "${color}${text}${NC}"
}

# 函数:显示欢迎信息
show_welcome() {
    clear
    print_color $BLUE "=========================================="
    print_color $BLUE "          Conda 自动安装脚本"
    print_color $BLUE "=========================================="
    echo
    print_color $GREEN "本脚本将帮助您安装 Conda 环境"
    print_color $GREEN "支持 AMD64 和 ARM64 架构"
    echo
}

# 函数:检测系统架构
detect_architecture() {
    local arch=$(uname -m)
    case $arch in
        x86_64)
            echo "amd64"
            ;;
        aarch64|arm64)
            echo "arm64"
            ;;
        *)
            echo "unknown"
            ;;
    esac
}

# 函数:获取用户输入
get_user_input() {
    # 检测系统架构
    local detected_arch=$(detect_architecture)
    
    print_color $YELLOW "检测到系统架构: $detected_arch"
    echo
    
    # 选择架构
    if [ "$detected_arch" != "unknown" ]; then
        print_color $YELLOW "请选择安装架构:"
        echo "1) 使用检测到的架构 [$detected_arch] (推荐)"
        echo "2) 手动选择架构"
        read -p "请输入选择 [1-2] (默认:1): " arch_choice
        
        case $arch_choice in
            1|"")
                ARCH=$detected_arch
                ;;
            2)
                print_color $YELLOW "请选择架构:"
                echo "1) AMD64 (x86_64) - 安装 Anaconda"
                echo "2) ARM64 (aarch64) - 安装 Miniconda"
                read -p "请输入选择 [1-2]: " manual_arch
                case $manual_arch in
                    1) ARCH="amd64" ;;
                    2) ARCH="arm64" ;;
                    *) 
                        print_color $RED "无效选择，使用检测到的架构"
                        ARCH=$detected_arch
                        ;;
                esac
                ;;
            *)
                print_color $RED "无效选择，使用检测到的架构"
                ARCH=$detected_arch
                ;;
        esac
    else
        print_color $YELLOW "无法自动检测架构，请手动选择:"
        echo "1) AMD64 (x86_64) - 安装 Anaconda"
        echo "2) ARM64 (aarch64) - 安装 Miniconda"
        read -p "请输入选择 [1-2]: " manual_arch
        case $manual_arch in
            1) ARCH="amd64" ;;
            2) ARCH="arm64" ;;
            *) 
                print_color $RED "无效选择，退出安装"
                exit 1
                ;;
        esac
    fi
    
    # 选择安装路径
    echo
    if [ "$ARCH" = "amd64" ]; then
        default_path="$HOME/anaconda3"
    else
        default_path="$HOME/miniconda3"
    fi
    
    print_color $YELLOW "请选择安装路径:"
    echo "1) 默认路径 [$default_path] (推荐)"
    echo "2) 自定义路径"
    read -p "请输入选择 [1-2] (默认:1): " path_choice
    
    case $path_choice in
        1|"")
            INSTALL_PATH=$default_path
            ;;
        2)
            read -p "请输入自定义安装路径: " custom_path
            if [ -n "$custom_path" ]; then
                INSTALL_PATH=$custom_path
            else
                print_color $RED "路径不能为空，使用默认路径"
                INSTALL_PATH=$default_path
            fi
            ;;
        *)
            print_color $RED "无效选择，使用默认路径"
            INSTALL_PATH=$default_path
            ;;
    esac
    
    # 确认安装
    echo
    print_color $YELLOW "安装摘要:"
    echo "架构: $ARCH"
    if [ "$ARCH" = "amd64" ]; then
        echo "版本: Anaconda"
    else
        echo "版本: Miniconda"
    fi
    echo "安装路径: $INSTALL_PATH"
    echo
    
    read -p "确认开始安装? [y/N]: " confirm_install
    if [[ ! $confirm_install =~ ^[Yy]$ ]]; then
        print_color $RED "安装已取消"
        exit 0
    fi
}

# 函数:安装conda
install_conda() {
    local arch=$1
    local install_path=$2
    local url=""
    
    if [ "$arch" = "amd64" ]; then
        url="https://repo.anaconda.com/archive/Anaconda3-2023.09-0-Linux-x86_64.sh"           
        print_color $BLUE "安装 Anaconda (AMD64) 到 $install_path"
    else
        url="https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-aarch64.sh"           
        print_color $BLUE "安装 Miniconda (ARM64) 到 $install_path"
    fi

    # 步骤1:检查并安装依赖
    print_color $GREEN "[1/5] 检查系统依赖..."
    sudo apt-get update -y > /dev/null
    sudo apt-get install -y wget bzip2 libgl1-mesa-glx > /dev/null

    # 步骤2:下载安装包
    print_color $GREEN "[2/5] 下载Conda..."
    wget -c -q --show-progress $url -O conda_installer.sh

    # 步骤3:静默安装
    print_color $GREEN "[3/5] 执行静默安装..."
    bash conda_installer.sh -b -p $install_path

    # 步骤4:环境配置
    print_color $GREEN "[4/5] 配置环境变量..."
    $install_path/bin/conda init bash
    echo "export PATH=\"$install_path/bin:\$PATH\"" >> ~/.bashrc
    
    # 步骤5:清理
    print_color $GREEN "[5/5] 清理安装文件..."
    rm conda_installer.sh
    
    # 显示完成信息
    echo
    print_color $GREEN "=========================================="
    print_color $GREEN "           安装完成!"
    print_color $GREEN "=========================================="
    echo
    print_color $BLUE "Conda 已成功安装到: $install_path"
    echo
    print_color $YELLOW "下一步操作:"
    echo "1. 重新打开终端"
    echo "2. 或运行: source ~/.bashrc"
    echo "3. 验证安装: conda --version"
    echo
    print_color $BLUE "感谢使用 Conda 自动安装脚本!"
}

# 主程序
main() {
    show_welcome
    get_user_input
    
    # 开始安装
    echo
    print_color $BLUE "开始安装 Conda..."
    install_conda $ARCH $INSTALL_PATH
}

# 检查是否以非交互模式运行
if [[ "$1" == "--non-interactive" ]]; then
    if [ $# -ne 3 ]; then
        echo "非交互模式用法: $0 --non-interactive [amd64|arm64] [安装路径]"
        exit 1
    fi
    ARCH=$2
    INSTALL_PATH=$3
    install_conda $ARCH $INSTALL_PATH
else
    # 运行主程序
    main
fi